#!/usr/bin/env python

# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import json
import os
import sys
from copy import deepcopy
from pathlib import Path
from typing import Dict, Tuple

import cv2
import numpy as np
import torch
import torchvision.transforms.functional as F  # type: ignore  # noqa: N812
from tqdm import tqdm  # type: ignore

from lerobot.common.datasets.lerobot_dataset import LeRobotDataset


def select_rect_roi(img):
    """
    Allows the user to draw a rectangular ROI on the image.

    The user must click and drag to draw the rectangle.
    - While dragging, the rectangle is dynamically drawn.
    - On mouse button release, the rectangle is fixed.
    - Press 'c' to confirm the selection.
    - Press 'r' to reset the selection.
    - Press ESC to cancel.

    Returns:
        A tuple (top, left, height, width) representing the rectangular ROI,
        or None if no valid ROI is selected.
    """
    # Create a working copy of the image
    clone = img.copy()
    working_img = clone.copy()

    roi = None  # Will store the final ROI as (top, left, height, width)
    drawing = False
    index_x, index_y = -1, -1  # Initial click coordinates

    def mouse_callback(event, x, y, flags, param):
        nonlocal index_x, index_y, drawing, roi, working_img

        if event == cv2.EVENT_LBUTTONDOWN:
            # Start drawing: record starting coordinates
            drawing = True
            index_x, index_y = x, y

        elif event == cv2.EVENT_MOUSEMOVE:
            if drawing:
                # Compute the top-left and bottom-right corners regardless of drag direction
                top = min(index_y, y)
                left = min(index_x, x)
                bottom = max(index_y, y)
                right = max(index_x, x)
                # Show a temporary image with the current rectangle drawn
                temp = working_img.copy()
                cv2.rectangle(temp, (left, top), (right, bottom), (0, 255, 0), 2)
                cv2.imshow("Select ROI", temp)

        elif event == cv2.EVENT_LBUTTONUP:
            # Finish drawing
            drawing = False
            top = min(index_y, y)
            left = min(index_x, x)
            bottom = max(index_y, y)
            right = max(index_x, x)
            height = bottom - top
            width = right - left
            roi = (top, left, height, width)  # (top, left, height, width)
            # Draw the final rectangle on the working image and display it
            working_img = clone.copy()
            cv2.rectangle(working_img, (left, top), (right, bottom), (0, 255, 0), 2)
            cv2.imshow("Select ROI", working_img)

    # Create the window and set the callback
    cv2.namedWindow("Select ROI")
    cv2.setMouseCallback("Select ROI", mouse_callback)
    cv2.imshow("Select ROI", working_img)

    print("Instructions for ROI selection:")
    print("  - Click and drag to draw a rectangular ROI.")
    print("  - Press 'c' to confirm the selection.")
    print("  - Press 'r' to reset and draw again.")
    print("  - Press ESC to cancel the selection.")

    # Wait until the user confirms with 'c', resets with 'r', or cancels with ESC
    while True:
        key = cv2.waitKey(1) & 0xFF
        # Confirm ROI if one has been drawn
        if key == ord("c") and roi is not None:
            break
        # Reset: clear the ROI and restore the original image
        elif key == ord("r"):
            working_img = clone.copy()
            roi = None
            cv2.imshow("Select ROI", working_img)
        # Cancel selection for this image
        elif key == 27:  # ESC key
            roi = None
            break

    cv2.destroyWindow("Select ROI")
    return roi


def select_square_roi_for_images(images: dict) -> dict:
    """
    For each image in the provided dictionary, open a window to allow the user
    to select a rectangular ROI. Returns a dictionary mapping each key to a tuple
    (top, left, height, width) representing the ROI.

    Parameters:
        images (dict): Dictionary where keys are identifiers and values are OpenCV images.

    Returns:
        dict: Mapping of image keys to the selected rectangular ROI.
    """
    selected_rois = {}

    for key, img in images.items():
        if img is None:
            print(f"Image for key '{key}' is None, skipping.")
            continue

        print(f"\nSelect rectangular ROI for image with key: '{key}'")
        roi = select_rect_roi(img)

        if roi is None:
            print(f"No valid ROI selected for '{key}'.")
        else:
            selected_rois[key] = roi
            print(f"ROI for '{key}': {roi}")

    return selected_rois


def get_image_from_lerobot_dataset(dataset: LeRobotDataset):
    """
    Find the first row in the dataset and extract the image in order to be used for the crop.
    """
    row = dataset[0]
    image_dict = {}
    for k in row:
        if "image" in k:
            image_dict[k] = deepcopy(row[k])
    return image_dict


def save_cropped_images(
    original_dataset: LeRobotDataset,
    crop_params_dict: Dict[str, Tuple[int, int, int, int]],
    output_dir: str,
    resize_size: Tuple[int, int] = (128, 128),
):
    """
    简化版函数，仅将裁剪后的图像保存到指定目录，不创建新的数据集。
    
    Args:
        original_dataset (LeRobotDataset): 原始数据集
        crop_params_dict (Dict[str, Tuple[int, int, int, int]]): 裁剪参数字典
        output_dir (str): 输出目录
        resize_size (Tuple[int, int]): 调整大小
    """
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    # 保存裁剪参数
    with open(output_path / "crop_params.json", "w") as f:
        json.dump(crop_params_dict, f, indent=4)
    
    # 创建图像输出目录
    for key in crop_params_dict.keys():
        key_dir = output_path / key.replace(".", "_")
        key_dir.mkdir(parents=True, exist_ok=True)
    
    # 处理每一帧
    for frame_idx in tqdm(range(len(original_dataset))):
        frame = original_dataset[frame_idx]
        episode_idx = frame["episode_index"].item()
        frame_idx_in_ep = frame["frame_index"].item()
        
        # 处理每个需要裁剪的图像
        for key in crop_params_dict:
            if key in frame:
                # 应用裁剪和调整大小
                top, left, height, width = crop_params_dict[key]
                value = frame[key]
                cropped = F.crop(value, top, left, height, width)
                resized = F.resize(cropped, resize_size)
                resized = resized.clamp(0, 1)
                
                # 将图像转换为numpy数组并保存
                img_np = (resized.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
                img_path = output_path / key.replace(".", "_") / f"ep_{episode_idx:06d}_frame_{frame_idx_in_ep:06d}.png"
                cv2.imwrite(str(img_path), cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Crop rectangular ROIs from a LeRobot dataset - 仅图像处理版本")
    parser.add_argument(
        "--repo-id",
        type=str,
        default="lerobot",
        help="The repository id of the LeRobot dataset to process.",
    )
    parser.add_argument(
        "--root",
        type=str,
        default=None,
        help="The root directory of the LeRobot dataset.",
    )
    parser.add_argument(
        "--crop-params-path",
        type=str,
        default=None,
        help="The path to the JSON file containing the ROIs.",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default=None,
        help="输出目录，默认为<root>_cropped_images",
    )
    args = parser.parse_args()

    # 加载数据集
    print(f"加载数据集: {args.repo_id} 从 {args.root}")
    dataset = LeRobotDataset(repo_id=args.repo_id, root=args.root)

    # 获取图像并转换为OpenCV格式
    images = get_image_from_lerobot_dataset(dataset)
    images = {k: v.cpu().permute(1, 2, 0).numpy() for k, v in images.items()}
    images = {k: (v * 255).astype("uint8") for k, v in images.items()}

    # 获取裁剪参数
    if args.crop_params_path is None:
        rois = select_square_roi_for_images(images)
    else:
        with open(args.crop_params_path) as f:
            rois = json.load(f)

    # 打印选择的ROI
    print("\n选择的矩形感兴趣区域 (top, left, height, width):")
    for key, roi in rois.items():
        print(f"{key}: {roi}")

    # 设置输出目录
    output_dir = args.output_dir if args.output_dir else str(Path(args.root).with_name(f"{Path(args.root).name}_cropped_images"))
    
    # 保存裁剪后的图像
    print(f"\n正在处理图像并保存到: {output_dir}")
    save_cropped_images(
        original_dataset=dataset,
        crop_params_dict=rois,
        output_dir=output_dir,
        resize_size=(128, 128),
    )
    
    print(f"\n处理完成! 裁剪后的图像已保存到: {output_dir}")
    print(f"裁剪参数已保存到: {Path(output_dir) / 'crop_params.json'}")
